import random
import subprocess
import os
import argparse
import shutil
from multiprocessing import Pool

def process_segment(segment_info):
    current_time, next_time, input_video_path, speed_factor, volume_factor, noise_factor, segment_file = segment_info
    high_freq = random.randint(*random.choice([(10, 80), (16000, 18000)]))
    cmd = [
        'ffmpeg',
        '-f', 'lavfi', 
        '-i', f'sine=frequency={high_freq}:duration={next_time-current_time}',
        '-ss', str(current_time),
        '-to', str(next_time),
        '-i', input_video_path,
        '-filter_complex',
        f'[1:v]setpts={1/speed_factor}*PTS[v];'  # 修正输入流索引为[1:v]
        f'[1:a]atempo={speed_factor},volume={volume_factor}[a1];'
        f'anoisesrc=c=1:a=\'{noise_factor}\'[a2];'
        f'[0:a]volume=1[high_freq];'
        f'[a1][high_freq][a2]amix=inputs=3:duration=first:dropout_transition=0[a]',  # 修正混音输入
        '-map', '[v]',
        '-map', '[a]',
        '-c:v', 'libx264',
        '-preset', 'ultrafast',
        '-c:a', 'aac',  # 显式指定音频编码器
        '-shortest',  # 确保输出与最短输入流对齐
        f"tmp/{segment_file}"
    ]
    subprocess.run(cmd)

def adjust_video_speed_randomly(input_video_path, output_video_path):
    # 创建 tmp 文件夹，如果不存在
    if not os.path.exists('tmp'):
        os.makedirs('tmp', exist_ok=True)

    # 获取视频时长
    cmd = f'ffprobe -v error -show_entries format=duration -of default=noprint_wrappers=1:nokey=1 "{input_video_path}"'
    result = subprocess.run(cmd, shell=True, capture_output=True, text=True)
    duration = float(result.stdout.strip())

    current_time = 0
    segment_files = []
    segment_index = 0
    segment_infos = []

    while current_time < duration:
        # 随机生成间隔时间，范围在 2 到 10 秒之间
        interval = random.uniform(2, 10)
        next_time = min(current_time + interval, duration)

        speed_factor = random.uniform(0.95, 1.15)
        volume_factor = random.uniform(0.6, 1.4)
        noise_factor = random.uniform(0, 0.01)

        segment_file = f"segment_{segment_index}.mp4"
        segment_files.append(segment_file)
        segment_infos.append((current_time, next_time, input_video_path, speed_factor, volume_factor, noise_factor, segment_file))

        current_time = next_time
        segment_index += 1

    # 使用多进程并行处理视频片段
    with Pool() as pool:
        pool.map(process_segment, segment_infos)

    # 生成用于合并片段的文件列表
    segments_txt_path = os.path.join('tmp', 'segments.txt')
    with open(segments_txt_path, 'w') as f:
        for segment in segment_files:
            f.write(f"file '{segment}'\n")

    # 合并所有片段
    merge_cmd = [
        'ffmpeg',
        # '-hwaccel', 'cuvid',
        # '-hwaccel_output_format', 'cuda',
        '-y',
        '-f', 'concat',
        '-safe', '0',
        '-i', segments_txt_path,
        '-c', 'copy',
        output_video_path
    ]
    subprocess.run(merge_cmd)

    # 删除临时文件
    shutil.rmtree('tmp')

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='随机调整视频速度和音量')
    parser.add_argument('input_video', type=str, help='输入视频文件路径')
    parser.add_argument('output_video', type=str, help='输出视频文件路径')
    args = parser.parse_args()

    adjust_video_speed_randomly(args.input_video, args.output_video)
    